package main

import "fmt"

func updateMatrix(mat [][]int) [][]int {
	res := make([][]int, 0)
	for i := 0; i < len(mat); i++ {
		res = append(res, make([]int, len(mat[i])))
	}
	for i, line := range mat {
		for j := 0; j < len(line); j++ {
			if mat[i][j] != 0 && res[i][j] == 0 {
				bfs(mat, i, j, res)
			}
		}
	}
	return res
}

// 单源BFS，大矩阵超时，根节点只有一个，例如树的遍历
func bfs(mat [][]int, row, col int, res [][]int) {
	if mat[row][col] == 0 {
		res[row][col] = 0
		return
	}
	var queue [][]int
	var round int
	queue = append(queue, []int{row, col})
	for len(queue) != 0 {
		size := len(queue)
		found := false
		for i := 0; i < size; i++ {
			head := queue[0]
			i := head[0]
			j := head[1]
			queue = queue[1:]
			if i-1 >= 0 {
				if mat[i-1][j] == 0 {
					found = true
					res[i-1][j] = 0
					break
				} else {
					queue = append(queue, []int{i - 1, j})
				}
			}
			if i+1 < len(mat) {
				if mat[i+1][j] == 0 {
					found = true
					res[i+1][j] = 0
					break
				} else {
					queue = append(queue, []int{i + 1, j})
				}
			}
			if j-1 >= 0 {
				if mat[i][j-1] == 0 {
					found = true
					res[i][j-1] = 0
					break
				} else {
					queue = append(queue, []int{i, j - 1})
				}
			}
			if j+1 < len(mat[i]) {
				if mat[i][j+1] == 0 {
					found = true
					res[i][j+1] = 0
					break
				} else {
					queue = append(queue, []int{i, j + 1})
				}
			}
		}
		round++
		if found {
			res[row][col] = round
			return
		}
	}
}

func main() {
	mat := [][]int{
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{1, 1, 1},
		{0, 0, 0},
	}
	fmt.Println(updateMatrix(mat))
}
